# plot_gdp_power_yoy.py
# 功能：读取GDP和用电量累计值，计算同比增速，绘制双Y轴交互图（下拉切换省份）
# Y轴颜色与折线一致：GDP蓝，用电量红

import os
import pandas as pd
import plotly.graph_objects as go
from plotly.subplots import make_subplots

# ------------------ 路径设置 ------------------
data_dir = os.path.join('.', 'data')
print(data_dir)

gdp_path = os.path.join(data_dir, 'GDP累计值.xlsx')
power_path = os.path.join(data_dir, '用电量累计值.xlsx')
output_html = os.path.join('../output', 'GDP累计同比与用电量累计同比混频关联分析.html')


# ------------------ 读取数据 ------------------
# GDP：跳过第一行，第二行为列名，第一列为日期
gdp_df = pd.read_excel(gdp_path, header=1, index_col=0)
gdp_df.index = pd.to_datetime(gdp_df.index)

# 用电量：第一行为列名，第一列为日期
power_df = pd.read_excel(power_path, header=0, index_col=0)
power_df.index = pd.to_datetime(power_df.index)

# ------------------ 提取省份 ------------------
def extract_province(col):
    return col.split(':')[0]

gdp_cols = {extract_province(col): col for col in gdp_df.columns}
power_cols = {extract_province(col): col for col in power_df.columns}

# 共有省份
common_provinces = sorted(set(gdp_cols.keys()) & set(power_cols.keys()))

if len(common_provinces) == 0:
    raise ValueError("没有共同的省市数据，请检查列名格式是否为 '省份:指标:类型'")

print(f"共找到 {len(common_provinces)} 个共有省市：{common_provinces}")

# ------------------ 计算同比函数 ------------------
def calculate_yoy(series):
    """
    计算同比增速：(当前值 - 去年同期值) / |去年同期值|
    返回单位为 % 的 Series
    """
    # 按年月对齐：只保留有去年同期数据的点
    series = series.dropna()
    yoy = []
    dates = []
    for date, value in series.items():
        last_year_date = date - pd.DateOffset(years=1)
        if last_year_date in series.index:
            last_year_value = series[last_year_date]
            if abs(last_year_value) > 0:  # 避免除以0
                growth = (value - last_year_value) / abs(last_year_value) * 100
                yoy.append(growth)
                dates.append(date)
    return pd.Series(yoy, index=dates).sort_index()

# ------------------ 创建图表 ------------------
fig = make_subplots(specs=[[{"secondary_y": True}]])

# 默认显示第一个省份
first_prov = common_provinces[0]
gdp_col = gdp_cols[first_prov]
power_col = power_cols[first_prov]

# 提取原始数据
gdp_series = gdp_df[gdp_col].dropna()
power_series = power_df[power_col].dropna()

# 计算同比
gdp_yoy = calculate_yoy(gdp_series)
power_yoy = calculate_yoy(power_series)

# 添加初始数据
fig.add_trace(
    go.Scatter(
        x=gdp_yoy.index,
        y=gdp_yoy,
        mode='lines+markers',
        name='GDP同比增速',
        line=dict(color='blue'),
        marker=dict(size=6),
        legendgroup='gdp',
        showlegend=True
    ),
    secondary_y=False,
)

fig.add_trace(
    go.Scatter(
        x=power_yoy.index,
        y=power_yoy,
        mode='lines+markers',
        name='用电量同比增速',
        line=dict(color='red'),
        marker=dict(size=4),
        legendgroup='power',
        showlegend=True
    ),
    secondary_y=True,
)

# ------------------ 构建下拉菜单 ------------------
dropdown_buttons = []

for province in common_provinces:
    gdp_col = gdp_cols[province]
    power_col = power_cols[province]

    gdp_series = gdp_df[gdp_col].dropna()
    power_series = power_df[power_col].dropna()

    gdp_yoy = calculate_yoy(gdp_series)
    power_yoy = calculate_yoy(power_series)

    idx = common_provinces.index(province)
    visibility = [False] * len(common_provinces) * 2
    visibility[idx * 2] = True
    visibility[idx * 2 + 1] = True

    button = dict(
        label=province,
        method='update',
        args=[
            {'visible': visibility},
            {'title': f"{province} GDP与用电量同比增速"}
        ]
    )
    dropdown_buttons.append(button)

# ------------------ 布局设置 ------------------
fig.update_layout(
    updatemenus=[
        dict(
            buttons=dropdown_buttons,
            direction='down',
            showactive=True,
            x=0.1,
            xanchor='left',
            y=1.15,
            yanchor='top',
            font=dict(size=12),
            bordercolor="#ccc",
            borderwidth=1,
            bgcolor="#f9f9f9"
        )
    ],
    title={
        'text': f'各省市GDP与用电量同比增速趋势图（当前：{first_prov}）',
        'x': 0.5,
        'xanchor': 'center',
        'y': 0.95,
        'font': {'size': 18}
    },
    xaxis_title="日期",
    hovermode="x unified",
    height=700,
    legend=dict(
        orientation="h",
        x=0.5, xanchor="center",
        y=1.08, yanchor="top"
    ),
    margin=dict(t=150),
    # 设置Y轴颜色
    yaxis=dict(
        title="GDP同比增速（%）",
        title_font=dict(color="blue"),
        tickfont=dict(color="blue"),
        zeroline=True,
        zerolinecolor='gray',
        zerolinewidth=1
    ),
    yaxis2=dict(
        title="用电量同比增速（%）",
        title_font=dict(color="red"),
        tickfont=dict(color="red"),
        zeroline=False  # 避免双零线重叠
    )
)

# ------------------ 添加其余省份的 trace（默认隐藏）------------------
for i, province in enumerate(common_provinces):
    if province == first_prov:
        continue
    gdp_col = gdp_cols[province]
    power_col = power_cols[province]

    gdp_series = gdp_df[gdp_col].dropna()
    power_series = power_df[power_col].dropna()

    gdp_yoy = calculate_yoy(gdp_series)
    power_yoy = calculate_yoy(power_series)

    fig.add_trace(
        go.Scatter(
            x=gdp_yoy.index,
            y=gdp_yoy,
            mode='lines+markers',
            name='GDP同比增速',
            line=dict(color='blue'),
            marker=dict(size=6),
            legendgroup='gdp',
            showlegend=False,
            visible=False
        ),
        secondary_y=False,
    )

    fig.add_trace(
        go.Scatter(
            x=power_yoy.index,
            y=power_yoy,
            mode='lines+markers',
            name='用电量同比增速',
            line=dict(color='red'),
            marker=dict(size=4),
            legendgroup='power',
            showlegend=False,
            visible=False
        ),
        secondary_y=True,
    )

# ------------------ 输出HTML ------------------
fig.write_html(
    output_html,
    config={'displayModeBar': True},
    include_plotlyjs='cdn',
    auto_open=False
)

print(f"✅ 已生成同比增速交互图表，保存至：{output_html}")
print("💡 用浏览器打开，使用顶部下拉菜单切换省份。")
print("📈 图中显示：(累计值今年 - 去年同期累计值) / |去年同期| × 100%")